
#include "mem.h"

using namespace AscendC;


template <typename T>
class TokenShiftB1T1Align
{
public:
    __aicore__ inline TokenShiftB1T1Align() {}
    __aicore__ inline void Init(
        GM_ADDR x, 
        GM_ADDR rwkvag, 
        GM_ADDR h0, 
        GM_ADDR xr, 
        GM_ADDR xw, 
        GM_ADDR xk,
        GM_ADDR xv, 
        GM_ADDR xa, 
        GM_ADDR xg, 
        GM_ADDR ht, 
        GM_ADDR workspace,
        uint32_t aivNum,
        uint32_t B_length,
        uint32_t T_length,
        uint32_t C_length,
        uint32_t tiling_size,
        uint32_t core_rows,
        uint32_t core_rows_remain,
        uint32_t X_length_align)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero");

        this->aivNum = aivNum;

        this->B_length = B_length;
        this->T_length = T_length;
        this->C_length = C_length;
        this->tiling_size = tiling_size;
        this->ALIGN_NUM = 32 / sizeof(T);
        uint32_t core_length = this->core_rows * C_length;
        uint32_t startpointer = GetBlockIdx() * C_length;

        x_global.SetGlobalBuffer((__gm__ T *)x, core_length);

        rwkvag_global.SetGlobalBuffer((__gm__ T *)rwkvag + startpointer, C_length);
        h0_global.SetGlobalBuffer((__gm__ T *)h0 , C_length);
        
        xr_global.SetGlobalBuffer((__gm__ T *)xr, core_length);   
        xw_global.SetGlobalBuffer((__gm__ T *)xw, core_length);   
        xk_global.SetGlobalBuffer((__gm__ T *)xk, core_length);   
        xv_global.SetGlobalBuffer((__gm__ T *)xv, core_length);   
        xa_global.SetGlobalBuffer((__gm__ T *)xa, core_length);   
        xg_global.SetGlobalBuffer((__gm__ T *)xg, core_length);        
        ht_global.SetGlobalBuffer((__gm__ T *)ht, C_length); 

        // AscendC::TEventID eventID = GetTPipePtr()->AllocEventID<AscendC::HardEvent::MTE2_V>(); 

        DataCopy(g_RWKVAG0Local, rwkvag_global, C_length); 
        // SetFlag<AscendC::HardEvent::MTE2_V>(eventID);
        // WaitFlag<AscendC::HardEvent::MTE2_V>(eventID); 
        
        DataCopy(xLocal, x_global, C_length); 
        // SetFlag<AscendC::HardEvent::MTE2_V>(eventID);
        // WaitFlag<AscendC::HardEvent::MTE2_V>(eventID); 

        DataCopy(h0Local, h0_global, C_length); 
        
        // SetFlag<AscendC::HardEvent::MTE2_V>(eventID);
        // WaitFlag<AscendC::HardEvent::MTE2_V>(eventID); 
        // GetTPipePtr()->ReleaseEventID<AscendC::HardEvent::MTE2_V>(eventID);
        PipeBarrier<PIPE_ALL>();

    }

    __aicore__ inline void Process() {

        Compute(0, this->tiling_size);
        CopyOutX(GetBlockIdx(), this->tiling_size);
               
    }
    
    __aicore__ inline void Compute(uint32_t startpos, uint32_t length) {

        Sub(h0Local, h0Local, xLocal, length);
        PipeBarrier<PIPE_V>();
        Mul(xrLocal, h0Local, g_RWKVAG0Local, length);
        PipeBarrier<PIPE_V>();
        Add(xrLocal, xrLocal, xLocal, length);
        PipeBarrier<PIPE_V>(); 

        if(GetBlockIdx() == 0) {
            DataCopy(ht_global, xLocal, length); 
        }    

    }    
 
    __aicore__ inline void CopyOutX(uint32_t blockidx, uint32_t length) {
        // LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        switch (blockidx)
        {
            case 0:
                DataCopy(xr_global, xrLocal, length); 
                break;
            case 1:
                DataCopy(xw_global, xrLocal, length); 
                break;
            case 2:
                DataCopy(xk_global, xrLocal, length); 
                break;
            case 3:
                DataCopy(xv_global, xrLocal, length); 
                break;  
            case 4:
                DataCopy(xa_global, xrLocal, length); 
                break;
            case 5:
                DataCopy(xg_global, xrLocal, length); 
                break;                                      
        }

        // PipeBarrier<PIPE_ALL>();

    }

private:
    GlobalTensor<T> x_global;
    GlobalTensor<T> rwkvag_global;
    GlobalTensor<T> h0_global;
    GlobalTensor<T> xr_global, xw_global, xk_global, xv_global, xa_global, xg_global, ht_global; 

    uint32_t B_length;
    uint32_t T_length;
    uint32_t C_length;
    uint32_t tiling_size;
    uint32_t core_rows;
    uint32_t core_rows_remain;
    uint32_t startPointer;
    uint32_t X_length_align;
    uint32_t ALIGN_NUM;
    uint32_t prevLinePointer;
    uint32_t aivNum;

    AsdopsBuffer<ArchType::ASCEND_V220> buf;

    AscendC::LocalTensor<T> g_RWKVAG0Local = buf.GetBuffer<BufferType::ASCEND_UB, T>(7680);
    AscendC::LocalTensor<T> xLocal = buf.GetBuffer<BufferType::ASCEND_UB, T>(15360);
    AscendC::LocalTensor<T> h0Local = buf.GetBuffer<BufferType::ASCEND_UB, T>(23040);
    AscendC::LocalTensor<T> xrLocal = buf.GetBuffer<BufferType::ASCEND_UB, T>(30720);

};